package cn.jbolt.ai.core.chain.node;

import cn.jbolt.ai.core.chain.node.action.aichat.AIChatNodeProvider;
import cn.jbolt.ai.core.chain.node.action.apicall.ApiCallNodeProvider;
import cn.jbolt.ai.core.chain.node.action.dbquery.DbQueryNodeProvider;
import cn.jbolt.ai.core.chain.node.action.funccall.FuncCallNodeProvider;
import cn.jbolt.ai.core.chain.node.action.funccall.FuncChatNodeProvider;
import cn.jbolt.ai.core.chain.node.action.intention.DefaultIntentionNodeProvider;
import cn.jbolt.ai.core.chain.node.action.klb.KlbChatNodeProvider;
import cn.jbolt.ai.core.chain.node.action.klb.KlbSearchNodeProvider;
import cn.jbolt.ai.core.chain.node.action.rewrite.RewriteNodeProvider;
import cn.jbolt.ai.core.chain.node.flow.varset.VarSetNodeProvider;
import cn.jbolt.ai.core.chain.node.output.toclient.OutputToClientNodeProvider;
import cn.jbolt.ai.core.chain.node.flow.ifcondition.IfConditionNodeProvider;
import cn.jbolt.ai.core.chain.node.provider.BooleanNodeProvider;
import cn.jbolt.ai.core.chain.node.provider.CommonNodeProvider;
import cn.jbolt.ai.core.chain.node.provider.EventNodeProvider;
import cn.jbolt.ai.core.chain.node.provider.IntentionNodeProvider;
import cn.jbolt.ai.core.chain.node.start.StartFromUserInputNode;
import cn.jbolt.util.SpringContextUtil;
import org.springframework.stereotype.Component;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 节点处理类注册中心
 */
@Component
public class NodeProviderCenter {


    private static final ConcurrentHashMap<String, Class> NODE_PROVIDER_HOLDER = new ConcurrentHashMap<>();

    static {
        //内置节点
        NODE_PROVIDER_HOLDER.put(InnerNodeType.START_FROM_USERINPUT, StartFromUserInputNode.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.AI_CHAT, AIChatNodeProvider.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.FUNC_CALL, FuncCallNodeProvider.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.FUNC_CHAT, FuncChatNodeProvider.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.KLB_SEARCH, KlbSearchNodeProvider.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.KLB_CHAT, KlbChatNodeProvider.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.AI_INTENTION, DefaultIntentionNodeProvider.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.AI_REWRITE, RewriteNodeProvider.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.OUTPUT_TO_CLIENT, OutputToClientNodeProvider.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.API_CALL, ApiCallNodeProvider.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.IF_CONDITION, IfConditionNodeProvider.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.DB_QUERY, DbQueryNodeProvider.class);
        NODE_PROVIDER_HOLDER.put(InnerNodeType.VAR_SET, VarSetNodeProvider.class);
    }

    private static void checkIsSpringBean(Class nodeProviderClass) {
        if (SpringContextUtil.getBean(nodeProviderClass) == null) {
            throw new IllegalArgumentException("节点类型["+nodeProviderClass+"]未注册为Spring Bean");
        }
    }

    /**
     * 添加事件节点
     * @param nodeName
     * @param nodeProviderClass
     * @param <T>
     */
    public static <T extends EventNodeProvider> void registerEventNodeProvider(String nodeName, Class<T> nodeProviderClass) {
        checkIsSpringBean(nodeProviderClass);
        NODE_PROVIDER_HOLDER.put(nodeName, nodeProviderClass);
    }


    /**
     * 添加布尔类型节点
     * @param nodeName
     * @param nodeProviderClass
     * @param <T>
     */
    public static <T extends BooleanNodeProvider> void registerBooleanNodeProvider(String nodeName, Class<T> nodeProviderClass) {
        checkIsSpringBean(nodeProviderClass);
        NODE_PROVIDER_HOLDER.put(nodeName, nodeProviderClass);
    }


    /**
     * 添加意图识别类型节点
     * @param nodeName
     * @param nodeProviderClass
     * @param <T>
     */
    public static <T extends IntentionNodeProvider> void registerIntensionNodeProvider(String nodeName, Class<T> nodeProviderClass) {
        checkIsSpringBean(nodeProviderClass);
        NODE_PROVIDER_HOLDER.put(nodeName, nodeProviderClass);
    }
    /**
     * 添加普通类型节点
     * @param nodeName
     * @param nodeProviderClass
     * @param <T>
     */
    public static <T extends CommonNodeProvider> void registerCommonNodeProvider(String nodeName, Class<T> nodeProviderClass) {
        checkIsSpringBean(nodeProviderClass);
        NODE_PROVIDER_HOLDER.put(nodeName, nodeProviderClass);
    }



    /**
     * 根据节点名称获取节点实例
     * @param nodeName 节点名称
     * @return 节点实例
     * @throws IllegalArgumentException 如果未注册该名称的节点类
     */
    public static Object getNodeProvider(String nodeName) {
        Class nodeClass = NODE_PROVIDER_HOLDER.get(nodeName);
        if (nodeClass == null) {
            throw new IllegalArgumentException("未注册的节点类型: " + nodeName);
        }
        return SpringContextUtil.getBean(nodeClass);
    }
}
